[Model] Enable ConvNeXt V2 with batch > 1 (issue #255)#295
Open
YWHyuk wants to merge 68 commits into
Open
Conversation
… feed Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime, loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in dma-fine-grained so multi-tile-K and conv loads do not collide; strip the reduction accum marker from the memory_barrier slot.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs, and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes; the trace logs the first DRAM response and the assigned systolic array, and scopes the compute barrier to its dispatch.
Default to the trace path; fix uninitialized Instruction fields, the matmul accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and dma_wait fusion; bound concurrent dispatches to the spad, round-robin work-items within a partition, benchmark autotune and run the multi-tenant scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose [P3-trace] logs; fix the key.mlir compile race in load().
… runtime model Replace the trace bridge's accumulated special cases with one dataflow rule and clean up the runtime that consumes it. Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all current writers (occupancy=ISSUE when both are systolic-array ops, else latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a matmul that reads and writes the same buffer = a commutative accumulator): skip its read edge and UNION its write, waiting only the non-matmul init seed and not ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector epilogue or the store wait every K matmul (fixes the pure-vector store that an empty COMPUTE_BAR let slip). Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA drain), so the store JOINs the matmul writers directly. The whole emit/loader chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a stale producer fails loudly instead of emitting records the bridge would drop. Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue). Model compute-output spad footprint in the SRAM version/capacity machinery so buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes from the DMA records that touch the same buffer (a buf_bytes pre-pass); an in-place buffer (accumulator, relu) is version-transparent so footprint is not double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP issue points (release before the COMP skip path so a skipped matmul still frees). Runtime: collapse child_inst / _pipeline_children into one event-indexed _deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot release queue and the async-load wakeup into one _due_events timed-effect table drained by process_due_events. Both are behavior-preserving (byte-identical). Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at init), and the round-robin disable mode is removed. Degenerate traces (a consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather than asserts that vanish under NDEBUG. Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in warns at command build.
… spad/2
The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace is self-documenting. Both are read from togsim_runtime.h at codegen time -- the version from the TOGSIM_ABI_VERSION define, the call-format text from a marked block next to the declarations -- so they never drift when the ABI changes.
…rint The bridge sums each work-item's distinct-buffer footprint (each buffer once, so a reduction's reloads of the same section do not inflate it) onto its Tile. can_issue then admits two concurrent dispatches only when each fits half the spad; a dispatch larger than spad/2 runs alone with the whole spad, so two work-items never compete for the shared spad and deadlock -- a runtime safety net beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are logged on the TILE_SCHEDULED trace line for debugging.
…election select_tile fed the epilogue buffer count to gemm_combination_mapping as max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers were already freed and reusable by the epilogue. The codegen lays every buffer out as a disjoint .spad global and never reuses freed space, so the estimate undercounted the real footprint: for a fused matmul+relu kernel the budget came to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088 B/lane including the relu output buffer. Under the full-spad guard this was harmless (89088 < 131072), but the spad/2 guard rejected it and crashed test_transformer_fusion with SpadOverflowError, since the heuristic path has no tile-shrink fallback. Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per epilogue node plus one per extra read operand, matching what the codegen emits. For the matmul+relu kernel the budget now equals the actual footprint exactly, and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2. Liveness-based spad reuse and broadcast over-allocation are tracked separately as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane scratchpad usage (spad_usage = stack_size + spad_size) and rejected the tiling when that exceeded the spad/2 budget. But only the .spad section actually lives in the scratchpad -- it is pinned there by the --section-start=.spad link option. The kernel stack is in main memory (sp is set up by pk in the -m region, not at the scratchpad vaddr), so it does not consume the scratchpad and must not be charged against it. The scalar frame is also shared across lanes, not per-lane, so adding it double-counted. On small configs (8x8) this falsely rejected feasible tilings: the wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to spare but were tripped over the spad/2 gate purely by the ~200-800 B stack term, crashing compile with SpadOverflowError. Drop the stack term; the .spad-only check still correctly rejects real buffer overflows (e.g. sparsity, which is fixed separately by f05ac8a).
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP padding entries are skipped instead of dereferencing uninitialized garbage, fixing the host store segfault on 8-lane (wrapper3) configs where a tile's split axis is not a multiple of vlane_stride*n_vu.
Add pytorchsim_functional_verify_per_kernel, a sub-option of pytorchsim_functional_mode that localizes the first compiled kernel whose Spike output diverges from a CPU reference. When enabled, the generated wrapper compares every realized buffer (the output of each fused kernel) against a CPU "golden" computed once per call by running the original aten graph (V.graph.module) on CPU with the same inputs, via an fx.Interpreter that records each node's output by name. A buffer is mapped to its originating fx node through V.graph.get_buffer().origin_node, so the check reports the kernel, the originating op, the offending indices and the max abs diff, then raises at the first divergence (stop-at-first). Comparison granularity is the fused-cluster output, the finest observable in a fused pipeline. Auto-disabled when functional mode is off (no Spike values to verify); the config accessor AND-gates the key with functional mode. Codegen bakes the verify_init/verify_check calls into the wrapper only when enabled at compile time, so clear the codegen cache when toggling. Tolerances via TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4). extension_functional_verify.py holds the graph registry and the runtime golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level; extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and dropped the prologue's extra-read operands entirely, so a softmax-fused attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The codegen lays the softmax max/sum operands out as their own disjoint weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so the estimate undercounted the real footprint: on the 32x32 wrapper2 config (16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted kernel used 16896 B/lane and overflowed the budget by 512 B, crashing test_transformer with SpadOverflowError. Under the 128x128 full-budget config the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the GEMM template (commit f05ac8a), now on the BMM prologue path. Add an n_prologue_extra_read knob to gemm_combination_mapping that charges each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N) .spad buffer, matching what the codegen emits, and have the BMM template count those operands the same way codegen does (the one numel-matching main input reuses the matmul-operand buffer; every other read gets its own global). Tile selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile (8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose check. The new parameter defaults to 0, leaving the GEMM and conv callers unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism, limitations, code map) and point to it from CLAUDE.md: a one-line entry in the debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump v1..v12 history was noise nobody reads.
[Frontend] Budget fused-prologue spad buffers in BMM tile selection
…s off format_dma_inst_issued_trace_line and format_instruction_detail_line are only ever used as the message argument to trace_instruction_line (a spdlog::trace sink). Because the argument is evaluated eagerly at the call site, the kernel builds a formatted std::string for every issued/finished instruction and then spdlog drops it whenever the level is above trace (the default). Bail out of both formatters when trace logging is disabled so the fmt::format work is skipped entirely. This removes wasted work but is not the simulation-speed bottleneck (that is the per-cycle ready-queue rescan in Core::cycle, addressed separately); the conv kernel wall time is unchanged within noise.
Core::cycle() walked every tile's ready-instruction list on every simulated cycle to find one instruction to issue. Profiling an 8x8 conv kernel showed this scan is ~96% of simulation time and, within it, the walk is >99%: the ready list holds ~2000 instructions (mostly blocked on the SRAM-capacity / weight-slot throttle), only one issues per cycle, and the blocked ones are never removed, so the same ~2000 are re-examined every cycle (~2.6k list iterations/cycle, ~1e9 over 400k cycles). At small arrays this dominates because tiny tiles inflate both the ready-list length and the number of DMA-wait stall cycles. Gate the scan behind a per-Core _issue_dirty flag. The scan's outcome can only change when the ready set grows or a resource frees, so set the flag on exactly those events: - ready set grows: Instruction::dec_ready_counter, on enqueueing a now-ready instruction, sets the OWNING core's _issue_dirty via _owner_dirty_ref (wired in Core::issue when the tile is dispatched) -- per-core so a remote core's enqueue does not force every core to rescan; - SRAM freed (release_sram) and weight slot freed (apply_due) set _issue_dirty; - a new dispatch (issue) and a successful issue keep it set. On a cycle with none of these the scan would re-walk the same blocked instructions and issue nothing, so it is skipped. Instructions are never moved between queues, so there is no re-admission churn. Issue-identical: it never changes which instruction issues or when, only whether the (side-effect-free under those conditions) scan runs. The 8x8 conv kernel produces the same 403026 cycles as before, with its TOGSim wall dropping from 214.5s to 41.9s (~5x, 71% of scans skipped). A forced-scan invariant check (no skipped cycle ever issues or inline-finishes a zero-cycle COMP) found 0 violations across 39 kernels spanning GEMM, BMM, softmax and conv. DMA-bound kernels (the ones that hit the 6h CI cap) stall more and gain more.
In Core::cycle()'s issue scan, a COMP with compute_cycle == 0 finishes inline and calls instructions.erase(it), but does not set `issued` or break -- so control falls through to the `it++` at the end of the loop body, incrementing the std::list iterator that erase() just invalidated. That reads a freed node and is undefined behavior; it usually limps along because the freed node's next pointer is briefly intact, but under a different allocator / sanitizer / build it can crash or skip-or-duplicate the remaining ready instructions in that tile, corrupting issue order. Use the iterator erase() returns and continue, the standard erase-in-loop idiom. No behavior change on the path that happened to work; it just removes the UB.
…X TOG The main compile/sim path no longer generates or selects the legacy ONNX Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv (the build-skip now keys on trace.so), and TOGSimulator.run_standalone always drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path (extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR). origins (which FX nodes a kernel came from) is preserved: logged per kernel run and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at the comment so the current parser is unaffected. Also drop the dead tog_file param from mlir_gem5_compile_command, migrate scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor addresses and --attributes_list is no longer a Simulator option), and refresh the CLAUDE.md TOG-generation notes.
Rework per-kernel codegen so the loop body is an ordered list of load->compute->store steps (class Step + push_step + codegen_loops iteration) instead of the fixed dma_loads/compute/dma_stores buffers. A step may be compute-only or transfer-only; the compute_idx loop is emitted only for steps that actually have compute content. Drop the ad-hoc self.masks buffer: the reduction-tail mask (get_mask) now emits into the compute buffer, since a mask is just compute. Use the step model to express indirect (gather/scatter) access cleanly. When the indirect index is produced in the same pass, convert_indirect_indexing pushes a new step for the offset build, so the index load, the offset build and the gather become separate steps bridged through spad. This replaces the compute_dependecy -> dma_stores hack in both convert_indirect_indexing and load(). push_step clones the CSE so the name counter is shared but the dedup cache is per-step (each step is its own compute_idx region). Validated: x[idx+2]+1 emits two steps and matches CPU; add/matmul/reduce/softmax/layernorm and gather/scatter unchanged.
…ptor
Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.
The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).
Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
…ride, compute-loop multi-dim offset Three refinements on top of the explicit offset descriptor: 1. Detect indirect access by an explicit symbol set instead of an "tmp" substring match. indirect_indexing now records str(index_var) in self.indirect_symbols; a _has_indirect(expr) helper tests the index free_symbols against it; the former "tmp"-string sites (store, get_dma_info, convert entry/indirect_dims/stride) use it. Removes the now-dead _find_indirect from lower_dma_to_gemmini. 2. Single indirect dim: pass the raw index spad and let the MVIN apply the gather stride per position (CONFIG4 offset_stride) instead of a vector_load/muli/vector_store round-trip that pre-multiplied the stride. emit_transfer carries offset_stride; decompose copies the attribute; lower programs CONFIG4 with it (was hardcoded to 1). 3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided indices, which a single CONFIG4 channel cannot do, so the sum stays in the kernel -- but build it in the compute loop (chunked by compute_vec_size, not a single tile-wide vector) and store it to a DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix) is not clobbered. push_step separates the offset-build loop from the gather that reads it. Adds multi-dim gather and index-reuse cases to test_indirect_access. Validated: indirect/scatter/embedding + the two new multi-dim cases pass; add/matmul/softmax regression 0.
…emmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash).
…-safe) _GoldenInterpreter recorded each node output on CPU but returned the original (npu) tensor to downstream nodes, and cast the recorded value to float32. A device-baked op (e.g. arange on npu) or a consumer that requires operand device agreement (aten.index needs its index tensors on the indexed tensor's device) then saw a mixed CPU/npu pair and the whole golden threw, silently disabling the per-kernel verify for that graph -- e.g. MobileNet's depthwise conv lowered to constant_pad + index gather. run_node now moves EVERY tensor output to CPU (preserving dtype) and returns it, so downstream nodes see CPU operands; only the recorded golden is cast to float32 for the allclose compare, so int index tensors stay usable.
The cat template's _calculate_input_tile_sizes handed each input a concat-dim tile of min(extent, remaining_spad_budget). When that budget did not cover an input's full concat extent, the tile width could fail to divide the extent (e.g. extent 64, tile 62 on the 8-lane config). The copy loop steps by that tile over [0, extent), so a non-dividing tile produces a ragged final iteration with no remainder mask: the MVIN over-reads past the input and the MVOUT over-writes past the output, corrupting the result. This surfaced on the 8x8 (wrapper3) CI config, where the per-input spad budget lands at 126 for two 64-wide inputs, truncating the second tile to 62. It scrambled ~45% of the cat output and broke LlamaDecoderLayer (rotary rotate_half cat) and the diffusion UNet skip-connection concat. Larger arrays give a bigger budget so both inputs got the full 64-wide (dividing) tile, which is why 32x32 and 128x128 passed. Snap each input's concat-dim tile down to the largest divisor of its extent that fits the budget, so every loop iteration is full-width and in-bounds. When the budget already covers the full extent the tile is unchanged, so 32x32/128x128 codegen is a no-op. Verified with a minimal cat repro (64+64 -> 128 on 8x8: fails before, diff 0 after) and the real test_llama LlamaDecoderLayer at 8x8 (now passes).
The vlane index_expr lowering reconstructs each lane element's logical index along the split axis as stride_dim + outer_dim * nr_vector_lane, where outer_dim is the sublane-row index. The next sublane row starts vector_lane * vlane_stride elements later (vector_lane lanes, each holding vlane_stride positions), but the multiplier used nr_vector_lane alone, dropping the vlane_stride factor. Any index tensor (e.g. arange/iota) whose split axis spills into a second sublane row then had every row past the first come out shifted, producing wrong values. This fired on the 8x8 (wrapper3) config: arange(32) at 8 lanes with vlane_stride 2 needs two sublane rows, so positions 16-31 came out as p-8. Larger arrays keep the 32-wide axis in a single row (outer_dim always 0), which is why 32x32/128x128 were unaffected. It broke test_llama's LlamaModel (rotary position ids) at 8x8. Multiply the outer/row term by vector_lane * vlane_stride. This equals the old value whenever vlane_stride == 1, so single-row splits and larger configs are a codegen no-op. Verified with a minimal arange(32) repro at 8x8 (positions 16-31 wrong before, diff 0 after) and the full test_llama LlamaModel at 8x8 (now passes).
init_tile_size gave a 1-D tile of 2*vlane_stride*vector_lane regardless of the data extent. When the extent is smaller than that tile the reduction/pointwise MVIN reads past the end of the DRAM buffer, and the compute loop folds the out-of-range elements into the result. For a sum reduction that means adjacent-DRAM garbage is squared/added into the sum. This broke test_single_perceptron's Loss on the 128x128 config: the MSE mean reduces 128 elements but the tile was 2*2*128 = 512, so the MVIN over-read 384 elements past the 128-wide input and roughly doubled the squared-error sum (loss diff ~4593). It fired on wide-lane configs (tile 512 > extent 128); 32x32 gives tile 128 == extent and 8x8 gives a tile that divides 128, so both were unaffected. Cap the 1-D tile to min(2*vlane_stride*vector_lane, extent) so the tile never exceeds the data. It is a no-op whenever extent >= tile, so larger reductions and other configs are unchanged. Verified with a minimal mse_loss(1-D) repro at 128x128 (over-reads before -> diff 0 after, across extents 64/100/127/200/300) and the real test_single_perceptron at 128x128 (Loss now passes).
torch.sort's frontend lowering exposes the sorted VALUES as the sole scheduler-visible template output; the INDICES buffer is written in place by the same kernel (make_inplace) but is not registered as an output. For argsort the values output is dead, so Inductor DCEs the whole sort kernel and the indices buffer is returned uninitialized (all-zeros), silently corrupting any argsort result (e.g. DeepSeek MoE expert routing). Register the in-place indices write as a MutationOutput owned by the sort operation so the scheduler keeps the kernel alive whenever indices is consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small MLIRTemplateBuffer subclass that can own extra MutationOutputs and have MLIRTemplateCaller.output_node build it; the sort lowering attaches the indices mutation via add_mutation_output. For templates with no mutation this is behaviorally identical to the base (get_outputs == [self]). Then gate the values MVOUT on the values output being live: in argsort it is pruned from the kernel args, so emitting its MVOUT referenced an undeclared %YV. The values scratchpad is still produced internally for the bitonic compare; only the dead DRAM store is skipped. Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort with both outputs live is unchanged; test_sort, test_matmul, test_cat, test_indirect_access unaffected by the output_node change. Note: x[argsort] end to end additionally needs the indirect-gather codegen fix (separate).
convert_indirect_indexing moves an indirect index into an offset spad and
zeros the indirect term in the DRAM base-offset expression, but it only
walked index.args. A BARE indirect index (x[idx], where the load index IS
the symbol so index.args is empty) was left unzeroed, so the DMA base
offset's affine.apply referenced the loop-local index SSA value before it
was defined ("use of undeclared SSA value name"). Any 1-D gather with a
runtime index tensor failed to compile at codegen; it was masked for
argsort only because the sort DCE bug left the indices all-zero (a
degenerate constant load).
Subs any remaining indirect symbol to 0 after the arg loop (mirrors the
store path). The per-position gather is carried by the offset spad, so the
base offset is correctly 0.
Verified: x[idx] (N=32/64) and x[x.argsort()] now compile and match CPU;
test_indirect_access.py (incl. scatter/index_add and multi-dim) and
test_sort.py unchanged.
_masked_bounds reverse-engineered per-dim padding from the single flat index offset and clamped each padded load to [pad, pad + input_extent). That recovery is ill-posed: the offset mixes the tap shift with the padding, and under channels_last the stride-1 channel axis absorbs the offset remainder, producing an impossible clamp (e.g. [4, 8) on a size-2 channel tile). Spike then zeroed the whole load, making channels_last depthwise conv 99.9 percent wrong while NCHW was fine. The clamp was also unnecessary: at pad positions the consumer already yields the correct value (a compute-side select for a padded gather, or the pad op's own fill), so reading OOB is harmless. Keep only the ragged tail clamp (glo=0, ghi=loop_extent), which is genuinely load-bearing for non-dividing tiles. Verified on 32x32: channels_last depthwise conv now matches CPU; NCHW conv and standalone ragged F.pad still pass. Broader regressions (padded reduction loads, full MobileNet) to be addressed separately.
codegen_epilogue_body() decided whether to emit the template buffer's MVOUT
by asking "did a fused epilogue already store something?":
if len(self.stores._lines) == 0:
template_store()
That is a proxy for the real question, "does anything outside this kernel still
read the template buffer?", and it is wrong in two of the four cases:
- epilogue + live template buffer -> store skipped, buffer never written
- no epilogue + dead template buffer -> store emitted for a pruned DRAM arg
The first case corrupts DeepSeek-V3's MoE gate. The gate is
sigmoid(mm(hidden, gate_w) + bias), so add+sigmoid fuse onto the GEMM template,
but the raw mm result is also read by a later gather kernel. Its buffer was
declared as a kernel output and never stored, so the gather read an
uninitialized buffer (all zeros), scrambling expert routing. Every config was
affected; only 8x8 crossed the test's loose atol=2e-1 (final Max abs diff
0.2546) while 32x32 landed just under it and passed.
Ask the real question instead. Inductor already answers it:
Kernel.remove_kernel_local_buffers() drops a buffer only when
Scheduler.can_buffer_be_removed_through_fusion() finds every user inside this
fused kernel. A buffer that survived removal therefore has an outside user and
must be stored. Record the template buffer's name in def_kernel() and
def_conv_kernel() (conv has its own implementation) and gate template_store()
on it, in both the main and the reduction_fusion branches. The proxy is gone.
Use the kernel-local self.removed_buffers, not V.graph.removed_buffers: the
former is what remove_buffer() populates, and the latter is still empty when
the store hook runs.
Note this makes a previously missing MVOUT issue, so DRAM traffic and cycle
counts rise for any model with a live template buffer under a fused epilogue.
That is the correct cost, not a regression, but reported cycles will shift.
Verified on the 8x8 config: test_deepseek_v3_base goes from Max abs diff 0.2546
to passing; a dead template buffer (relu(mm)) still prunes its DRAM arg and
emits a single MVOUT; the trace/TOGSim timing path runs with the extra store.
Regressions pass: test_bmm_reduction, test_matmul, test_bmm, test_conv2d,
test_cnn, test_group_conv, test_pool, test_reduce, test_softmax, test_layernorm,
test_mlp.
Implement abs, sign, isnan, isinf, fmod, bitwise shifts, copysign, erfc, hypot, cosh, sinh, acos, asin, atan, atan2, acosh, asinh, atanh, and log/ atan (VCIX) in the Inductor MLIR backend. Float-only transcendentals cast non-float inputs to f32 while leaving native float widths (f16/f64) untouched.
Port the C++ MathLog/MathAtanToVCIX patterns to the in-process Python VCIX pass: add math.log (opcode 2, imm 2) and math.atan (opcode 2, imm 3) to _MATH_VIV and MARKERS. Encoding matches the spike rs1 / gem5 VS1 sub-opcodes (sin=0, cos=1, log=2, atan=3). No C++ mlir-opt pass needed.
Add the pointwise op test under tests/ops/elementwise and register it in the test workflow allowlist.
Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
reduction_init returned "-inf"/"inf" for max/min regardless of dtype, so an
integer max/min reduction fed an invalid inf identity into every consumer:
the accumulator init, the template reduction init, and the masked-DMA tail
fill (torch.tensor(inf, dtype=int64) overflowed -- surfaced compiling swinv2
with batch>1). reduction_combine_vec also hardcoded the float multi_reduction
kinds <maximumf>/<minimumf>, which reject integer vectors.
Return the dtype's representable extreme from reduction_init for integer
max/min (max identity -inf -> iinfo.min, min identity +inf -> iinfo.max;
bool -> 0/1), and emit the signed-integer multi_reduction kinds <maxsi>/<minsi>
for integer element types. sum/prod need no change: the <add>/<mul> combining
kinds are polymorphic and ops.add/ops.mul already select arith.add{i,f} by
dtype (as ops.maximum/minimum already select maxsi vs maximumf).
Verified: int amax/amin/sum/prod reductions compile and match CPU; float
test_reduce (ReduceMax) and test_softmax unchanged. Also clears the swinv2
batch>1 masked-fill overflow (it then hits a separate unsupported
shifted-window view, tracked separately).
…lar DMA index
torch.roll has no Inductor lowering; it decomposes to index_select(fmod(...)) -- a
cyclic-shift gather whose ModularIndexing DMA index the affine-only descriptor
cannot express ("Unlinearized floor/mod in DMA index"). Rewrite it as the two
contiguous halves the shift produces, stitched by cat:
roll(x, s, d) == cat([slice(x, N-s, N), slice(x, 0, N-s)], d)
each half is a plain strided slice with no modulo. Do it as a lowering (removing
roll from the decomp table so it is not decomposed first) rather than a
decomposition, so we can also REALIZE the input with copy_input: roll's producer
is often a reshaped view (e.g. swinv2's window_reverse), and the slice offset would
otherwise fuse into that reshape's modular index, re-creating a shifted (p+c)%m
form. A plain .contiguous()/clone does NOT create the buffer boundary (Inductor
inlines it); copy_input does.
Verified: torch.roll (1-D / multi-dim / arbitrary shift) and roll+reduction match
CPU; swinv2 batch>1 now compiles through the shifted-window attention (previously
blocked by the roll's cyclic shift and then by the fused slice offset).
A composition of aligned reshapes on one iteration variable (e.g. swinv2's window partition fused with a later reshape) leaves a nested index like ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor simplify_with_ranges reduces. collect_boundaries then skips its cut points (the inner base is not a bare variable) and the affine-only DMA check rejects it. Add a general, pattern-free canonicalizer (flatten_nested_floormod / _as_digit): any nesting of FloorDiv/ModularIndexing over a single symbol is a digit extractor (v // A) % M and collapses to one level by composing divisors, from four divisibility-guarded algebraic identities. Applied in collect_boundaries and in _fold_with_ranges. Multi-variable inner arguments (e.g. a roll shift v+c) are left untouched. Verified: numeric equivalence on the swinv2 window index; test_reduce, test_layernorm, test_cat, test_indirect_access unaffected.
The wrapper header imported empty_strided but not empty_strided_cpu, so a graph that allocates a CPU-side buffer (e.g. swinv2's shifted-window attention mask, which Inductor keeps on CPU) failed at runtime with "NameError: name 'empty_strided_cpu' is not defined". Add the same binding the default Inductor PythonWrapperCodegen header provides.
…yout A per-lane reduction is lowered as a 2-D [reduction | batch] collapse: multi_reduction reshapes the accumulator to <reduction_numel x red_size> and reduces axis 0, which assumes the reduction axis is the outermost in-lane run and every non-reduction (batch) axis is inner. get_dma_info reorders the tile axes for reduction loads to guarantee that, but only for the 2-D and 3-D cases. The 4-D case mis-tested is_reduction (reduction_depth < 3, false once there are three batch dims) and then raised NotImplementedError, and rank 5+ fell through to the default row-major order -- both leaving the reduction axis innermost, so the reduce collapsed a batch axis instead of the reduction axis. This is hit whenever a non-reduction dim-merge is blocked and an extra batch axis stays in-lane: SwinV2 cosine window attention adds a gathered relative-position bias (table[idx[q,k], head]) before the softmax amax, keeping head separate from query (4-D tile, head bled into the max), and its post-attention LayerNorm var_mean splits the token axis into several loop vars (6-D tile, mean off by ~0.005). Generalize the reorder to any tile rank: a tile that carries the reduction axis places that axis-group outermost via range(r, L) + range(r-1, -1, -1), which reduces to the existing 3-D order for L=3. Adds test_reduce_gather_bias (fails max abs diff ~3.3 before the fix).
SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA index" because the shifted-window path (torch.roll composed with window partition/reverse) produced views axis-split could not linearize. With the roll->slice+cat lowering, the nested/shifted mod handling, and the reduction-axis layout fix, the whole backbone now compiles and matches CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2). Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to the other model tests, and list SwinV2 in the README model coverage.
conv_multi_tile_mapping (used for batch > 1 convs) collects every tile that fits SPAD into tile_candidates and returns them sorted, but it guarded the "no mapping" error on max_used_spad_size instead of on tile_candidates. max_used_spad_size is only bumped when a candidate also satisfies max_k_h_w <= k_h, and max_k_h_w is initialized to K_W, so it only ever fires for the full-kernel (k_h == K_H) tile. When that tile overflows SPAD -- e.g. a CLIP ViT-B/32 patch conv (3->768, 32x32 stride 32) at batch > 1 -- the guard raised "Cannot find a valid mapping" even though smaller-k_h tiles fit and were already in tile_candidates. The mapping variable it tracks is never returned. Guard on `not tile_candidates` instead, so the conv is rejected only when no tile fits at all. The returned candidate list is unchanged, so convs that already worked are unaffected. Adds a batched patch-conv case to test_conv2d.py. Fixes issue #252.
With the conv-mapping fix, CLIP's vision transformer runs with batch > 1 and matches CPU end to end (max diff ~1e-5). Add tests/models/test_clip.py (CLIPVisionModel, batch 2, patch 32), wire it as a self-hosted CI job, and list CLIP in the README model coverage.
ConvNextV2 with batch > 1 died in codegen with "Index names length
mismatch: 2 != 3". The scheduler had approved fusing a reduction epilogue
into a GEMM whose 2-D (M rows x N cols) frame cannot express it, and
set_ranges then found more loop ranges than the template's coordinate map
has entries.
It approved it because the eligibility check read the reduction's stride
out of the node's repr:
stride = [i.strip()[:-1].split(",")[-1].strip()
for i in str(node).split("\n") if "r0" in i][1]
The intent (reject a reduction over the contiguous, stride-1 axis, which
the GEMM reduction template cannot codegen) was right, but picking the
second line that mentions "r0" lands on the wrong index expression once
the node has more than two dimensions -- so ConvNextV2's channels-first
LayerNorm mean was let through. Reading the stride off the LoopBody's
index expression instead makes it decline for the right reason. Note a
MemoryDep's index cannot be used here: it is normalized to a flat
contiguous access and no longer carries the per-axis strides.
Fixing that exposed the deeper problem: only a template knows its output
coordinate frame, yet can_fuse_horizontal hardcoded a case per template
kind and re-guessed that frame. It also mutated nodes (revert_group) from
inside a predicate the scheduler calls speculatively for every candidate
pair, carried a dead isinstance(MaxPoolTemplate) special case, and
declined without ever saying why.
So the decision moves to the templates, following how Inductor's own
CUTLASS and CPP backends do it:
- MLIRTemplate.try_fuse_epilogue / try_fuse_prologue own every condition
that depends on the frame -- including the config gates -- and return a
FusionPlan or None. They are pure; a node mutation the fusion needs is
deferred into the plan's remap and applied from MLIRScheduling.fuse(),
the commit hook Inductor calls once it really fuses. The plan is
recomputed there rather than cached, since can_fuse runs many times per
pair and node identities change as fusions land.
- The scheduler keeps only graph facts: which node carries the template,
the direction, and that this is a single-node to single-node fusion.
Case 1 (pointwise) and Case 2 (reduction) collapse into one delegation.
- Every decline logs a reason, like Inductor's WhyNoFuseNames. Declining
is a normal answer: upstream declines reduction epilogues outright.
- A template declares REDUCTION_EPILOGUE_ALIASING, the coordinate map its
reduction-epilogue codegen addresses. render() consumes it, and its
length is exactly what set_ranges asserts against -- so the fusion gate
is that assert, raised as a decline instead of a crash. Templates
without one cannot absorb a reduction epilogue, which retires the
support_reduction_fusion flag. GEMM's map has two entries, BMM's three.
When the rows do not fit, LoopBody.merge_loops() (which returns a new
body, so this stays pure) says whether merging them would; if it would,
the merge becomes the plan's remap.
tests/ops/fusion covers pointwise, reduction, prologue, conv and
attention fusion and is unchanged, as are the kernels each of them fuses.
Adds a matmul whose reduction is over the contiguous axis: it must not
fuse, and wrongly fusing it returns wrong values rather than failing to
compile.
ConvNextV2 batch > 1 no longer hits the assertion; it now stops at an
unrelated SPAD overflow, so issue #255's report is not fully closed.
SpadOverflowError was raised with no arguments and the only diagnostic was a logger.debug line that is off by default, so a failing compile said nothing beyond "SPAD overflow occurred." and gave no way to tell which kernel, which tiling, or by how much. load() already has the measured usage, the budget, the tile size and the kernel's origins in scope. Put them in the exception message.
Spad globals follow a fixed convention: the MLIR memref carries the full tile shape (all lanes), the .spad C header that Spike links declares the per-lane slice, and the gem5 header declares the full tile. codegen_spad_buffer() emits exactly that, tile_numel_per_lane and tile_numel_per_lane * vector_lane. index_expr() had the two headers swapped. The iota scratch buffer it allocates for vectorized index arithmetic declared compute_vec_size * vector_lane entries in the .spad header and compute_vec_size entries in the gem5 header. So the buffer Spike sees is vector_lane times too large, and the one gem5 sees is one element too small: the initializer stores two elements per iteration, so its last iteration writes one slot past compute_vec_size. The oversize side is what breaks compiles. Since the spad guard tightened to spad/2 for double buffering, a kernel with compute_vec_size 64 spends 64 KB of its 64 KB budget on an iota that only ever uses 64 entries. ConvNextV2 with batch > 1 fails there: its channels-first LayerNorm var_mean kernel needs 81984 bytes/lane, of which 65536 is the iota and only 16448 is real data. Declare compute_vec_size + 1 entries per lane, and vector_lane times that for gem5. The var_mean kernel now measures 16968 bytes/lane.
Add tests/models/test_convnextv2.py, a self-hosted CI job for it, and a model coverage row. batch=2 with depths=[1,1,1,1] and image_size=64 now matches CPU end to end, max abs diff 5.0e-06. The test keeps the runner on self-hosted like the other model tests: the depthwise 7x7 conv lowers to one gather kernel per tap, so the graph is wide.
d9131ba to
1ff2ebb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #255. ConvNeXt V2 with
batch > 1died at codegen withAssertionError: Index names length mismatch: 2 != 3. Two independent frontend bugs were stacked behind that message; the second only became visible once the first was fixed.batch=2now matches CPU end to end, max abs diff 5.0e-06.1. The fusion gate read the reduction stride out of a repr string
can_fuse_horizontaldecided whether a reduction epilogue could be folded into a GEMM by parsing the node'srepr():The intent was right: a reduction over the output's contiguous (stride-1) axis cannot be expressed in the GEMM template's 2-D reduction-epilogue frame, so it must not fuse. But taking the second line containing
"r0"lands on the wrong index expression once the node has more than two dimensions. ConvNeXt V2's channels-first LayerNorm mean is exactly that shape, so it was let through, andset_ranges(mlir_common.py:584) assertedlen(dim_aliasing) == len(ranges)->2 != 3.The fix: templates own their own fusion conditions
Rather than repair the string parse, the decision layer now follows the pattern both upstream Inductor backends use (
CUDACPPScheduling._can_fuse_epilogue_implfor CUTLASS,template_fusion_with_epilogues_supportedfor CPP):can_fusestays pure, expressibility is judged from the IR, and declining is cheap and logged.MLIRTemplate.try_fuse_epilogue(template_node, existing_epilogues, node_to_fuse)andtry_fuse_prologue(...)own every frame-dependent condition, config gates included, and returnFusionPlan | None. They are pure, so the scheduler can call them speculatively.FusionPlan.remap, a zero-arg thunk applied fromMLIRScheduling.fuse()-- the commit hookScheduler.fuse_nodes_oncealready calls._single_template_epilogue_pair,_single_prologue_template_pair) and delegates. It loses 155 lines.why_no_fuse()and logs a reason at DEBUG onPyTorchSimFrontend.mlir.mlir_template.The stride check now reads the real index expression via
LoopBody.get_all_read_expr(buffer_name)and the real reduction symbol viabody.vars[1]. Note thatMemoryDep.indexcannot be used here: it is normalized to a flat contiguous access and carries no per-axis strides, soindex.coeff(reduce_var)is always 0, which silently disables the check and produces wrong values rather than a crash. That is presumably why the original parsed the repr.REDUCTION_EPILOGUE_ALIASINGThe coordinate map that
render()already used is hoisted to a class attribute. Its length is exactly whatset_rangesasserts against, so the fusion gate is that assert, raised early as a decline instead of late as a crash.Nonemeans the template cannot absorb a reduction epilogue at all, which retires thesupport_reduction_fusionflag. GEMM declares 2 entries (row + reduced N), BMM declares 3 (batch + row + reduced N) -- that length is the only difference between them, so the conditions live in the base class andmlir_gemm_template.pyends up with no fusion logic at all.When the epilogue needs more loop ranges than the frame addresses, the template asks
body.merge_loops()(which is pure, unlikeSchedulerNode.merge_loops()) whether collapsing would make it fit; if so the merge becomes theplan.remap. For ConvNeXt V2 it correctly refuses -- the reduced C axis sits between batch and spatial in the channels-first layout, so the rows are not contiguous -- and the reduction runs as its own kernel.2. The index-expression scratch buffer was sized per tile, not per lane
With the LayerNorm reduction now correctly running as its own kernel, ConvNeXt V2 hit
SpadOverflowError.Spad globals follow a fixed convention, visible in
codegen_spad_buffer()(mlir_codegen_backend.py:1552-1553):memref.globaltile_desc.get_mlir_shape().spadC header (Spike links this)tile_size // vector_lanetile_sizeindex_expr()had the two headers swapped. The iota buffer ([0, 1, ..., compute_vec_size-1]) it allocates for vectorized index arithmetic declaredcompute_vec_size * vector_laneentries in the.spadheader andcompute_vec_sizein the gem5 header. So Spike's buffer isvector_lanetimes too large, and gem5's is one element too small -- the initializer stores two elements per iteration (ops.broadcast(init_iter, 2), kept wide to avoid scalar ops), so its last iteration writes one slot pastcompute_vec_size.Only
[0, compute_vec_size)is ever touched: oneaffine.parallel (%i) = (0) to (compute_vec_size)writes it, oneaffine.vector_load %buf[0] : vector<compute_vec_size x index>reads it.Since the spad guard tightened to
spad/2for double buffering, the oversize side is fatal. ConvNeXt V2'svar_meankernel:buf0..buf3(data tiles)buf4, buf5(mean/var out)index_expr_64_spad[8192]The real data is 16 KB; an iota that uses 64 entries (512 B) exhausts the whole budget. The heuristic mapping path has no tile-shrink fallback --
SpadOverflowErroris only caught insideautotune()-- so this is a hard failure.The fix declares
compute_vec_size + 1entries per lane andvector_lanetimes that for gem5. The+1is the two-wide initializer store, made explicit in a comment; it is the same slack idea asmax(tile_numel_per_lane, 2)incodegen_spad_buffer(). Thevar_meankernel now measures 16968 bytes/lane, and every otherindex_exprkernel shrinks by the same factor (index_expr_16_spad: 2048 entries -> 17).The MLIR
memreftype is left alone: it already carries the full-tile count, matching the convention, and the.lldeclares the symbolexternalso the C header is what defines the storage.3.
SpadOverflowErrornow says which kernel overflowedIt was raised with no arguments and the only diagnostic was an off-by-default
logger.debug, so a failing compile printedSPAD overflow occurred.and nothing else.load()already has the measured usage, the budget, the tile size and the kernel's origins in scope. That is how the kernel above was identified.4. Test, CI, coverage
tests/models/test_convnextv2.py(batch=2,depths=[1,1,1,1],image_size=64), a self-hosted CI job, and a README coverage row.Verification
Test Passed, max abs diff5.0068e-06.conv2d(x, w, b, padding=3, groups=96)), the smallest op that emitsindex_expr: functional mode, rebuilt binaries, max abs diff7.6e-06. This exercises the resized buffer end to end and rules out an alignment regression from the smaller allocation.tests/ops/fusion/:test_matmul_reduction,test_bmm_reduction,test_matmul_activation,test_conv_fusion,test_prologue_fusionall still compile with their expected fusedorigins(e.g.mm, max_1andbmm, max_1). Worth watching on review: BMM also supports reduction fusion, so putting the reduction conditions in the GEMM template alone silently splitstest_bmm_reductioninto two kernels.tests/ops/fusion/test_matmul_reduction.py::test_matmul_reduce_last_dim: a matmul whose reduction is over the contiguous axis must not fuse. Fusing it anyway returns wrong values rather than failing to compile, so the test checks numbers. Every pre-existing test reducesdim=-2(a non-contiguous axis), so this hole existed.The original ConvNeXt V2 crash has no self-contained repro --
matmul+permute+mean,+addmm, and a full channels-first LayerNorm all fail to reproduce it, because the trigger needs a node with more than two dims whose second"r0"line is the wrong one. Rather than fabricate one, the model test guards it.Follow-ups (not in this PR)
try_fuse_epiloguealready takesexisting_epilogues, so enabling it only needs the scheduler's single-node gate relaxed (upstream CUTLASS does this).support_epilogue_fusion/support_prologue_fusionremain as declarative gates inside the template layer.